import numpy as np
import scipy.sparse as sp
from scipy.sparse.csgraph import shortest_path
from scipy.io import loadmat
import pickle
import time
import os
from tqdm import tqdm
from scipy.sparse.csgraph import connected_components

DATADIR = "data/"

def compute_shortest_path_matrix(dataset_name):
    """使用SciPy优化实现的最短路径计算"""
    # 加载原始MAT数据
    mat_path = os.path.join(DATADIR, f"{dataset_name}.mat")
    mat_data = loadmat(mat_path)

    # 获取原始稀疏矩阵并转换为无向图
    sparse_matrix = mat_data['homo']
    sparse_matrix = sparse_matrix.maximum(sparse_matrix.T)  # 确保矩阵对称

    # 转换为CSR格式并移除对角线元素（可选）
    sparse_matrix = sparse_matrix.tocsr()
    sparse_matrix.setdiag(0)  # 移除自环边

    # 计算最短路径
    print(f"Starting shortest path computation for {dataset_name}...")
    start_time = time.time()

    # 使用Dijkstra算法（自动处理非加权图）
    dist_matrix = shortest_path(sparse_matrix,
                                directed=False,
                                method='D',
                                unweighted=True)

    print(f"Computation completed in {time.time() - start_time:.2f} seconds")

    # 处理不可达节点
    dist_matrix[np.isinf(dist_matrix)] = -1

    # 转换为整数类型
    dist_matrix = dist_matrix.astype(np.int32)

    # 保存为pkl文件
    output_path = os.path.join(DATADIR, f"{dataset_name}_shortest_distance.pkl")
    with open(output_path, 'wb') as f:
        pickle.dump({
            'dist_matrix': dist_matrix,
            'sparse_matrix': sparse_matrix,
            'node_count': sparse_matrix.shape[0],
            'dataset': dataset_name
        }, f, protocol=pickle.HIGHEST_PROTOCOL)

    print(f"Saved shortest distance matrix for {dataset_name} at {output_path}")
    return dist_matrix


def test_shortest_path(dataset_name):
    """测试模块"""
    # 加载数据
    with open(os.path.join(DATADIR, f"{dataset_name}_shortest_distance.pkl"), 'rb') as f:
        data = pickle.load(f)

    dist_matrix = data['dist_matrix']
    n = data['node_count']

    # 测试1：对角线是否为0
    assert np.all(np.diag(dist_matrix) == 0), "Diagonal elements should be zero"

    # 测试2：矩阵对称性（无向图）
    assert np.allclose(dist_matrix, dist_matrix.T), "Matrix should be symmetric"

    # 测试3：随机抽样验证
    for _ in range(10):
        i, j = np.random.randint(0, n, 2)
    if dist_matrix[i, j] != -1:
    # 验证三角不等式
        k = np.random.randint(0, n)
    if dist_matrix[i, k] != -1 and dist_matrix[k, j] != -1:
        assert dist_matrix[i, j] <= dist_matrix[i, k] + dist_matrix[k, j], \
            f"Triangle inequality violated at ({i}, {j}, {k})"

    # 测试4：不可达节点验证
    unreachable_count = np.sum(dist_matrix == -1)
    print(f"Unreachable node pairs: {unreachable_count} ({(unreachable_count / (n * n)) * 100:.4f}%)")

    print(f"All tests passed for {dataset_name}")


def print_examples(dataset_name, num_samples=5):
    """打印最短距离矩阵的示例结果"""
    with open(os.path.join(DATADIR, f"{dataset_name}_shortest_distance.pkl"), 'rb') as f:
        data = pickle.load(f)

    dist_matrix = data['dist_matrix']
    sparse_matrix = data['sparse_matrix']

    n_components, labels = connected_components(sparse_matrix)
    component_size = np.sum(labels == labels[0])

    n = data['node_count']

    print(f"\n=== {dataset_name}数据集的最短距离矩阵 ===")
    print(f"总节点数: {n}")
    print(f"矩阵形状: {dist_matrix.shape}")
    print(f"最小距离: {np.min(dist_matrix[dist_matrix != -1])}")
    print(f"最大有效距离: {np.max(dist_matrix[dist_matrix != -1])}")

    # n_components, labels = connected_components(dist_matrix, directed=False)
    print(f"连通组件数量: {n_components}")
    # 查看节点0所属组件的大小
    # component_id = labels[0]
    # component_size = np.sum(labels == component_id)
    print(f"节点0所在组件包含 {component_size} 个节点")


    # 随机采样验证
    print(f"\n随机采样{num_samples}个节点对的距离：")
    np.random.seed(42)
    for _ in range(num_samples):
        i, j = np.random.randint(0, n, 2)
        distance = dist_matrix[i, j]
        if distance == -1:
            print(f"节点 {i:5d} 到节点 {j:5d} | 不可达")
        else:
            print(f"节点 {i:5d} 到节点 {j:5d} | 最短距离: {distance:2d} 跳")

    # 节点0的邻居分布统计
    sample_node = 0
    distances = dist_matrix[sample_node]
    valid_distances = distances[distances != -1]

    print(f"\n节点 {sample_node} 的邻居统计:")
    print(f"直接邻居数量（1跳）   : {np.sum(valid_distances == 1)}")
    print(f"间接邻居数量（2-3跳）: {np.sum((valid_distances >= 2) & (valid_distances <= 3))}")
    print(f"远距离邻居（>3跳）   : {np.sum(valid_distances > 3)}")
    unreachable = max(0, np.sum(distances == -1) - 1)
    print(f"不可达节点数量      : {unreachable}")

    # 矩阵片段展示
    print("\n矩阵片段（前5行前5列）:")
    print(dist_matrix[:5, :5].astype(int))

def main():
    for dataset in ["YelpChi", "Amazon"]:
        dist_matrix = compute_shortest_path_matrix(dataset)
        test_shortest_path(dataset)
        print_examples(dataset)

if __name__ == "__main__":
    main()